import os
import numpy as np
import h5py
import lightgbm as lgb
from y import Y

def load_dataset(path):
    h5_file = h5py.File(path, 'r')
    x = h5_file['x'][...].astype(np.float32)
    y = h5_file['y'][...][:, Y].astype(np.float32)
    return (x, y)

df_train = load_dataset('out/feature/train.h5')
df_val = load_dataset('out/feature/val.h5')
df_test = load_dataset('out/feature/test.h5')
train = lgb.Dataset(df_train[0], label=df_train[1], params={'verbose':-1})
val = lgb.Dataset(df_val[0], label=df_val[1], params={'verbose':-1})
test = lgb.Dataset(df_test[0], label=df_test[1], params={'verbose':-1})

param = dict(
    lambda_l1=0.05,
    objective='regression',
    metric='regression',
    bagging_fraction=0.4,
    feature_fraction=0.8,
    verbose=-1,
)
bst = lgb.train(param, train, num_boost_round=1000, valid_sets=[test], early_stopping_rounds=5, verbose_eval=False)
bst.save_model(f'out/models/lgb_{Y}.txt')

y = df_test[1]
y_pred = bst.predict(df_test[0])

r2 = 1 - (((y - y_pred)** 2) / ((y - y.mean())** 2).sum()).sum()
l2 = ((y-y_pred)**2).sum() / len(y)
print('r2', r2)
print('l2',l2)

